package net.sushine.dc.mongodbsync;

/**
 * @author zhang
 * @version 1.0
 * @date 2024-12-09 16:22
 */
import com.mongodb.BasicDBObject;
import com.mongodb.DBCollection;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import com.mongodb.WriteConcern;
import com.mongodb.client.AggregateIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.MongoIterable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import net.sushine.dc.mongodbsync.constant.Constants;
import net.sushine.dc.mongodbsync.util.MongoDBUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.operators.Keys;
import org.apache.flink.api.java.utils.ParameterTool;
import org.bson.Document;

public class MongoDBSync {
    public static void main(String[] args) {
        ParameterTool parameterTool = ParameterTool.fromArgs(args);
        String sourceHost = parameterTool.get("sourceHost");
        int sourcePort = Integer.parseInt(parameterTool.get("sourcePort"));
        String sourceUsername = parameterTool.get("sourceUsername");
        String sourcePassword = parameterTool.get("sourcePassword");
        String sourceDatabase = parameterTool.get("sourceDatabase");
        String sinkHost = parameterTool.get("sinkHost");
        int sinkPort = Integer.parseInt(parameterTool.get("sinkPort"));
        String sinkUsername = parameterTool.get("sinkUsername");
        String sinkPassword = parameterTool.get("sinkPassword");
        String sinkDatabase = parameterTool.get("sinkDatabase");
        String orgNo = parameterTool.get("orgNo");
        int pageSize = parameterTool.getInt("pageSize", 10000);
        if (StringUtils.isBlank(sourceDatabase) || StringUtils.isBlank(sinkDatabase) || StringUtils.isBlank(orgNo)) {
            System.out.println("miss required config:sourceDatabase,sinkDatabase,orgNo");
            System.exit(1);
        }
        long total = 0;
        try {
            MongoClient sourceClient = getClient(sourceHost, sourcePort, sourceUsername, sourcePassword);
            try {
                MongoClient sinkClient = getClient(sinkHost, sinkPort, sinkUsername, sinkPassword);
                MongoDatabase sourceMGDS = sourceClient.getDatabase(sourceDatabase);
                MongoIterable<String> sourceCollections = sourceMGDS.listCollectionNames();
                MongoDatabase sinkMGDS = sinkClient.getDatabase(sinkDatabase);
                MongoCollection<Document> sinkCollection = sinkMGDS.getCollection(Constants.COLLECTION_PREFIX + orgNo);
                Map<String, Integer> unitIdCountMap = new HashMap<>();
                AggregateIterable<Document> unitIdCountIterable = sinkCollection.aggregate(MongoDBUtil.count()).allowDiskUse(true);
                MongoCursor<Document> it = unitIdCountIterable.iterator();
                while (it.hasNext()) {
                    Document document = it.next();
                    unitIdCountMap.put(document.getString(DBCollection.ID_FIELD_NAME), document.getInteger("count"));
                }
                MongoCursor<String> it2 = sourceCollections.iterator();
                while (it2.hasNext()) {
                    String currentSCName = it2.next();
                    if (currentSCName.contains(Constants.COLLECTION_PREFIX)) {
                        MongoCollection<Document> sourceSC = sourceMGDS.getCollection(currentSCName);
                        long sourceCount = sourceSC.countDocuments();
                        System.out.println("collection:" + currentSCName + ",数据条数:" + sourceCount);
                        String unitId = currentSCName.split(Keys.ExpressionKeys.SELECT_ALL_CHAR_SCALA)[2];
                        if (unitIdCountMap.containsKey(unitId) && unitIdCountMap.get(unitId).intValue() > 0) {
                            long sinkCount = unitIdCountMap.get(unitId).intValue();
                            if (sourceCount != sinkCount) {
                                List<BasicDBObject> maxTS = MongoDBUtil.getMaxTS(unitId);
                                long maxTimestamp = sinkCollection.aggregate(maxTS).allowDiskUse(true).first().getLong("timestamp").longValue();
                                int sourceMatchesCount = sourceSC.aggregate(MongoDBUtil.countSource(maxTimestamp)).first().getInteger("count").intValue();
                                if (maxTimestamp == -1) {
                                    sourceMatchesCount = 0;
                                }
                                long pageCount = (sourceMatchesCount / pageSize) + 1;
                                for (int i = 1; i <= pageCount; i++) {
                                    List<BasicDBObject> pageQuery = MongoDBUtil.getPageQuery(((long) i) == pageCount ? sourceMatchesCount % pageSize : pageSize, pageSize, i, maxTimestamp, unitId);
                                    total = insertBatch(total, sinkCollection, sourceSC, pageQuery);
                                }
                            }
                        } else {
                            long pageCount2 = (sourceCount / pageSize) + 1;
                            for (int i2 = 1; i2 <= pageCount2; i2++) {
                                List<BasicDBObject> pageQuery2 = MongoDBUtil.getPageQuery(((long) i2) == pageCount2 ? ((int) sourceCount) % pageSize : pageSize, pageSize, i2, -1L, orgNo);
                                total = insertBatch(total, sinkCollection, sourceSC, pageQuery2);
                            }
                        }
                    }
                }
                closeClient(sinkClient);
                closeClient(sourceClient);
                System.out.println("此次同步了" + total + "条数据");
            } catch (Exception e) {
                System.out.println("sink mongodb establish connection failed:" + e.getMessage());
                throw e;
            }
        } catch (Exception e2) {
            System.out.println("source mongodb establish connection failed:" + e2.getMessage());
            throw e2;
        }
    }

    private static long insertBatch(long total, MongoCollection<Document> sinkCollection, MongoCollection<Document> sourceSC, List<BasicDBObject> pageQuery) {
        AggregateIterable<Document> pageResult = sourceSC.aggregate(pageQuery).allowDiskUse(true);
        List<Document> sinkDCS = new ArrayList<>();
        MongoCursor<Document> it = pageResult.iterator();
        while (it.hasNext()) {
            Document sourceDC = it.next();
            sinkDCS.add(extract(sourceDC));
        }
        sinkCollection.insertMany(sinkDCS);
        return total + sinkDCS.size();
    }

    private static Document extract(Document sourceDC) {
        Document document = new Document();
        document.put("unitId", (Object) sourceDC.getString("unitId"));
        document.put("timestamp", (Object) sourceDC.getLong("timestamp"));
        document.put("alertValue", (Object) sourceDC.getInteger("alertValue"));
        document.put("env", (Object) sourceDC.getString("env"));
        document.put("jobUnit", (Object) sourceDC.getString("jobUnit"));
        return document;
    }

    public static MongoClient getClient(String host, int port, String username, String password) {
        MongoClient mongoClient;
        if (StringUtils.isBlank(username)) {
            mongoClient = new MongoClient(host, port);
        } else {
            MongoCredential credential = MongoCredential.createCredential(username, "admin", password.toCharArray());
            mongoClient = new MongoClient(new ServerAddress(host, port), credential, MongoClientOptions.builder().writeConcern(WriteConcern.UNACKNOWLEDGED).build());
        }
        return mongoClient;
    }

    public static void closeClient(MongoClient mongoClient) {
        if (mongoClient != null) {
            mongoClient.close();
        }
    }

}
